#!/usr/bin/env python3
"""
Author: Haoran Peng
Email: gavinsweden@gmail.com
"""

from typing import Dict

import numpy as np

from .agent import Agent
from .constraints import Constraints


class CTNode:
    def __init__(self, constraints: Constraints, solution: Dict[Agent, np.ndarray]):
        self.constraints = constraints
        self.solution = solution
        self.cost = self.sic(solution)
        # solution是不同长度的，将solution补齐到最长的长度，用最后一个位置填充
        # self.solution = self.pad_solution(solution)

    # Sum-of-Individual-Costs heuristics
    @staticmethod
    def sic(solution):
        return sum(len(sol) for sol in solution.items())

    def __lt__(self, other):
        return self.cost < other.cost

    def __str__(self):
        return str(self.constraints.agent_constraints)

    def pad_solution(self, solution: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
        # 找到最长的路径长度
        max_length = max(len(path) for path in solution.values())
        # 遍历 solution，将每个路径补齐到最长的长度
        padded_solution = {}
        for agent, path in solution.items():
            if len(path) < max_length:
                # # 使用路径的最后一个位置进行填充
                # padding = np.full((max_length - len(path),) + path.shape[1:], path[-1])
                # padded_path = np.concatenate((path, padding), axis=0)
                # 使用从最后一个位置不断累加30的值进行填充
                last_value = path[-1]
                padding = np.array(
                    [
                        [last_value[0] + (i + 1) * 20, last_value[1]]
                        for i in range(max_length - len(path))
                    ]
                )
                padded_path = np.concatenate((path, padding), axis=0)
            else:
                padded_path = path
            padded_solution[agent] = padded_path

        return padded_solution
